# -*- coding:utf8 -*-
# @Time : 2023/1/31 11:39
# @Author : WanJie Wu

import torch.nn as nn
from transformers import BertModel


class BertClfModel(nn.Module):
    def __init__(self, model_name_or_path, num_classes):
        super(BertClfModel, self).__init__()
        self.bert = BertModel.from_pretrained(model_name_or_path)
        for param in self.bert.parameters():
            param.requires_grad = False
        self.fc = nn.Linear(self.bert.config.hidden_size, num_classes)

    def forward(self, input_ids, attention_mask, token_type_ids):
        output = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        out = self.fc(output.pooler_output)
        return out


if __name__ == "__main__":
    BertClfModel(
        model_name_or_path="/data/sdv1/wuwanjie/transformers/bert-base-chinese",
        num_classes=15
    )
